import os
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.chains import RetrievalQA
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.embeddings.dashscope import DashScopeEmbeddings
# 解决Pydantic版本警告问题
try:
    from pydantic import BaseModel
except ImportError:
    from langchain_core.pydantic_v1 import BaseModel

# 使用正确的文件路径（文件在当前目录下）
loader = PyMuPDFLoader("deepseek-v3-1-4.pdf")
pages = loader.load_and_split()

# 文档切分
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=200,
    length_function=len,
    add_start_index=True,
)
texts = text_splitter.create_documents(
    [page.page_content for page in pages[:1]]
)

embeddings = DashScopeEmbeddings(
    model="text-embedding-v1",
    dashscope_api_key="sk-2c036d80d14a443f8769e072bb09a2e1",
)
db = FAISS.from_documents(texts, embeddings)

# 修复方法名错误：get_retriever -> as_retriever
retriever = db.as_retriever(search_kwargs={"k": 5})

docs = retriever.invoke("deepseek v3有多少参数")

for doc in docs:
    print(doc.page_content)
    print("----")